import os
import matplotlib.pyplot as plt
import geopandas as gpd
import numpy as np
from matplotlib.font_manager import FontProperties
# 使用 macOS 的中文字体，比如 "PingFang SC"
font_prop = FontProperties(fname="/System/Library/Fonts/Hiragino Sans GB.ttc")

def plot_residual_map(df, save_path="residual_map.png", shp_path=None):
    if shp_path is None:
        # 获取项目根目录的绝对路径
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        shp_path = os.path.join(project_root, "utils/陕西省地市.shp")

    # 读取陕西省地市的 shp 文件
    try:
        gdf = gpd.read_file(shp_path)
    except Exception as e:
        print(f"读取 shp 文件时出错: {e}")
        return

    # 创建包含两个子图的画布
    fig, axes = plt.subplots(1, 2, figsize=(20, 8))

    # 定义异常等级及其对应的颜色
    anomaly_levels = df["异常等级"].unique()
    # colors = plt.cm.get_cmap('tab10', len(anomaly_levels))
    # level_color_map = {level: colors(i) for i, level in enumerate(anomaly_levels)}
    # 手动定义颜色，可根据异常等级数量调整，示例使用十六进制颜色码
    custom_colors = {
        "明确异常": "#FF0000",  # 红色
        "可疑异常": "#FFA500",  # 橙色
        "插值不确定": "#FFFF00",  # 黄色
        "IDW不准": "#0000FF",  # 蓝色
        "正常": "#008000"  # 绿色
    }
    # 确保自定义颜色覆盖所有异常等级
    level_color_map = {level: custom_colors.get(level, "#808080") for level in anomaly_levels}

    # 绘制陕西省地图边界线到两个子图
    for ax in axes:
        gdf.boundary.plot(ax=ax, edgecolor='black', linewidth=1)

    # 绘制站点克里金残差散点图
    scatter = axes[0].scatter(df["lon"], df["lat"], c=df["Kriging_residual"], cmap="RdBu", s=20, edgecolor='k')
    
    # 添加颜色条并设置 label 字体
    cbar = fig.colorbar(scatter, ax=axes[0])
    cbar.set_label("Kriging残差(mm)", fontproperties=font_prop)

    axes[0].set_xlabel("经度", fontproperties=font_prop)
    axes[0].set_ylabel("纬度", fontproperties=font_prop)
    axes[0].set_title("站点Kriging残差分布图", fontproperties=font_prop)
    axes[0].grid(True)

    # 绘制站点异常等级散点图
    for level in anomaly_levels:
        level_df = df[df["异常等级"] == level]
        axes[1].scatter(level_df["lon"], level_df["lat"], c=[level_color_map[level]], label=level, s=20, edgecolor='k')
    axes[1].legend(title="异常等级", loc='upper left', bbox_to_anchor=(1.1, 1), prop=font_prop, title_fontproperties=font_prop )
    axes[1].set_xlabel("经度", fontproperties=font_prop)
    axes[1].set_ylabel("纬度", fontproperties=font_prop)
    axes[1].set_title("站点异常等级分布图", fontproperties=font_prop)
    axes[1].grid(True)

    plt.tight_layout()
    # plt.savefig(save_path, dpi=300)
    plt.savefig(save_path, dpi=100)
    plt.close()
